#include <cmath>
#include <iostream>
#include <algorithm>

#include "agent.h"

using namespace std;

void agent::insert(solution* theSolution, solution* parent1, solution* parent2)
{
    thePool->insertSolution(theSolution);

    cout << theName << ": (";
    if (parent1) 
        cout << "#" << parent1->getNumber() << "," << parent1->getTrainFitness();
    if (parent2) 
        cout << "x" << "#" << parent2->getNumber() << "," <<  parent2->getTrainFitness();
    if (parent1) cout << "->";
    cout << "#" << theSolution->getNumber() << "," << theSolution->getTrainFitness() << ")";

    if (thePool->getBestSolution() == theSolution) cout << " BEST";

    cout << "\n";

    cout.flush();


}

double agent::calcEntropy(int left[], int right[])
{
    int leftSum = left[0] + left[1] + left[2] + left[3] + left[4];
    int rightSum = right[0] + right[1] + right[2] + right[3] + right[4];

    double entropy = 0.0;

    // Add entropy contribution from left child of split
    if (leftSum > 0)
    {
        for (int i = 0; i < 5; i++)
        {
            double p_i = (double) (left[i]) / leftSum;
            if (p_i > 0) entropy += leftSum / (double) (leftSum + rightSum) *
                                    -p_i * log(p_i)/log(2.0);
        }
    }
    // Add entropy contribution from right child of split
    if (rightSum > 0)
    {
        for (int i = 0; i < 5; i++)
        {
            double p_i = right[i] / (double) rightSum;
            if (p_i > 0) entropy += rightSum /
                                    (double) (leftSum + rightSum) * -p_i * log(p_i)/log(2.0);
        }
    }
    return entropy;
}


class FieldValueSorter
{
public:
  bool operator() (const pair<bool, double>& lh, const pair<bool, double>& rh)
  {  return lh.second < rh.second; }
};

node* agent::determineSplitContinuous(bitset<TRAIN_RECORD_COUNT>* includeRecs, int nFldNo)
{
    // Vector of FldValue
    vector< double > splits;


    // For roulette wheel:
    vector<double> slotSizes;
    double sum = 0;


    // Sort records, based on field value
    int totalCount[5] = {0,0,0,0,0};
    int totalCountSum = 0;
    // Pair of Target Value and Field Value
    vector<pair<int, double> > records;
  
    for (int i = 0; i < TRAIN_RECORD_COUNT; i++)
    {
        if ((*includeRecs)[i])
        {
            int type = thePool->getTrainingSet()->getNthRecord(i)->getResult();
            double fldVal = thePool->getTrainingSet()->getNthRecord(i)->getFieldValue(nFldNo);
            totalCount[type]++;
            totalCountSum++;

            records.push_back(make_pair(type, fldVal));
        } 
    }

    int empty[] = {0,0,0,0,0};
    double entropyParent = calcEntropy(totalCount,empty);

    if (entropyParent == 0.0) return NULL;

    // Sort records
    sort(records.begin(), records.end(), FieldValueSorter());

    // Find valid splits
    int currCountSum = 0;
    int currCount[5] = {0,0,0,0,0};
    double splitValue;
    if (records.size() > 0) splitValue  = records.begin()->second;
    for (vector<pair<int, double> >::iterator it = records.begin(); it != records.end(); ++it)
    {
        if (it->second != splitValue)
        {
            assert(it->second > splitValue);
            if (currCountSum >= MIN_SPLIT_SIZE &&
                totalCountSum - currCountSum >= MIN_SPLIT_SIZE)
            {

                int right[] = {totalCount[0] - currCount[0], totalCount[1] - currCount[1],
                               totalCount[2] - currCount[2], totalCount[3] - currCount[3],
                               totalCount[4] - currCount[4]};
                double entropy = calcEntropy(currCount, right);

                double slotSize = entropyParent - entropy;
                if (slotSize > 0)
                {
                    sum += slotSize;
    
                    // Insert a split between the last and current value
                    splits.push_back((splitValue + it->second) / 2.0);
                    slotSizes.push_back(slotSize);
                }
            }      
        }
        currCount[it->first]++;
        currCountSum++;

        splitValue = it->second;
    }
    node* retVal = NULL;

    if (slotSizes.size() == 0) return NULL;

    // Choose a slot
    double threshold = (double) rand() / RAND_MAX * sum;

    double nFldVal = *(splits.begin());
    for (int j = 0; threshold > 0 && j < splits.size(); ++j)
    {
        nFldVal = splits[j];
        threshold -= slotSizes[j];
    }

    retVal = new node(nFldNo, nFldVal);

    return retVal;
}

node* agent::determineSplitSymbolic(bitset<TRAIN_RECORD_COUNT>* includeRecs, int nFldNo)
{
    // Vector of FldValue
    vector< double > splits;

    // For roulette wheel:
    vector<double> slotSizes;
    double sum = 0;

    // Total counts for each type
    int totalCount[5] = {0,0,0,0,0};
    int totalCountSum = 0;

    // A map of Field Value to occurences of each type
    int occurences[72][5];
    for (int i = 0; i < 72; i++)
    {
        for (int j = 0; j < 5; j++)
            occurences[i][j] = 0;
    }

    int maxFldVal = -1;
    for (int i = 0; i < TRAIN_RECORD_COUNT; i++)
    {
        if ((*includeRecs)[i])
        {
            int type = thePool->getTrainingSet()->getNthRecord(i)->getResult();
            int fldVal = (int) thePool->getTrainingSet()->getNthRecord(i)->getFieldValue(nFldNo);
            if (fldVal > maxFldVal) maxFldVal = fldVal;

            totalCount[type]++;
            totalCountSum++;

            assert (fldVal < 72 && fldVal >= 0);
            occurences[fldVal][type]++;
        } 
    }

    int empty[] = {0,0,0,0,0};
    double entropyParent = calcEntropy(totalCount, empty);

    if (entropyParent == 0.0) return NULL;

    // Calculate the entropy for each field value
    vector<pair<int,double> > slots;
    for (int v = 0; v <= maxFldVal; v++)
    {
        int right[] = {totalCount[0] - occurences[v][0], totalCount[1] - occurences[v][1],
                       totalCount[2] - occurences[v][2], totalCount[3] - occurences[v][3],
                       totalCount[4] - occurences[v][4]};
        double entropy = calcEntropy(occurences[v],right);

        double slotSize = entropyParent - entropy;
        assert(slotSize >= 0);
        if (slotSize > 0)
        {
            sum += slotSize;
            slots.push_back(make_pair(v, slotSize));
        }
   
    }

    // Pick a random split
    if (slots.size() == 0) return NULL;

    // Choose a slot
    double threshold = (double) rand() / RAND_MAX * sum;

    int nFldVal = slots.begin()->first;
    double currEntropy = slots.begin()->second;
    for (int j = 0; threshold > 0 && j < slots.size(); ++j)
    {
        nFldVal = slots[j].first;
        currEntropy = slots[j].second;
        threshold -= slots[j].second;
    }

    set<int> values;
    values.insert(nFldVal);
    // See if adding a field value would reduce entropy
    while(values.size() < maxFldVal)
    {
        double bestEntropy = currEntropy;
        int bestNextFldVal = -1;
        int currCount[5] = {0,0,0,0,0};
        int currCountSum = 0;

        for (set<int>::iterator it = values.begin(); it != values.end(); ++it)
        {
            for (int i = 0; i < 5; i++)
            {
                currCount[i] += occurences[*it][i];
                currCountSum += occurences[*it][i];
            }
        }

        for (int i = 0; i <= maxFldVal; i++)
        {
            if (values.find(i) == values.end())
            {
                int left[] = {currCount[0] + occurences[i][0], currCount[1] + occurences[i][1],
                              currCount[2] + occurences[i][2], currCount[3] + occurences[i][3],
                              currCount[4] + occurences[i][4]};
                int right[] = {totalCount[0] - left[0], totalCount[1] - left[1],
                               totalCount[2] - left[2], totalCount[3] - left[3],
                               totalCount[4] - left[4]};
                double entropy = calcEntropy(left, right);
    
                if (entropy < bestEntropy)
                {
                    bestNextFldVal = i;
                    bestEntropy = entropy;
                }
            }
        }
        
        if (bestNextFldVal == -1) break;
        values.insert(bestNextFldVal);
    }

    // Make sure that we have valid split sizes
    int leftSize = 0; 
    for (set<int>::iterator it = values.begin(); it != values.end(); ++it)
    {
        for (int i = 0 ; i < 5; i++)
            leftSize += occurences[*it][i];
    }

    if (leftSize < MIN_SPLIT_SIZE || totalCountSum - leftSize < MIN_SPLIT_SIZE) return NULL;
    
    return new node(nFldNo, values);
}

node* agent::determineSplit(bitset<TRAIN_RECORD_COUNT>* includeRecs, int nFldNo)
{
    // Choose a field at random 
    if (nFldNo == -1)
    {
        nFldNo  = rand() % (MAX_FIELD - 1);
    }

    if (nFldNo == 1 || nFldNo == 2 || nFldNo == 3)
        return determineSplitSymbolic(includeRecs, nFldNo);
    else
        return determineSplitContinuous(includeRecs, nFldNo);
}

